/*-
 * Copyright (c) 2012, Columbia University
 * All rights reserved.
 *
 * This software was developed by Georgios Kontaxis <kontaxis@cs.columbia.edu>
 * at Columbia University, New York, NY, USA, in January 2012.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 *   * Redistributions of source code must retain the above copyright
 *     notice, this list of conditions and the following disclaimer.
 *   * Redistributions in binary form must reproduce the above copyright
 *     notice, this list of conditions and the following disclaimer in the
 *     documentation and/or other materials provided with the distribution.
 *   * Neither the name of Columbia University nor the
 *     names of its contributors may be used to endorse or promote products
 *     derived from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h> 
#include <sys/socket.h>
#include <netinet/in.h>
#include <netdb.h>
#include <unistd.h>
#include <netdb.h>
#include <signal.h>
#include <openssl/evp.h>

#include "config.h"
#include "cbuf.h"
#include "debug.h"
#include "splitter.h"
#include "secure_rw.h"
#include "auth.h"

/* number of messages buffered per connection 
 (in case of failure to deliver in real-time,
 e.g., due to disconnected endpoint) */
#define MSG_BUFFER_SIZE 10

/* array holding descriptors to the sockets listening 
 for incoming clients (one per configuration, i.e.,
 one per port to listen to) */
int * client_listeners = NULL;

/* array holding pointers to connection structure instances 
 for client connections (one per configuration) */
struct connection ** client_connections = NULL;

/* array holding pointers to connection strucutre instances 
 for agent connections (one per agent).
 agents of the same configuration are grouped
 together in a linked list with the primary
 agent being at the top of the list */
struct connection ** agent_connections = NULL;

/* count active configurations */
unsigned int config_count = 0;

void error(char *msg)
{
	perror(msg);
	exit(0);
}

void warning(char *msg)
{
	perror(msg);
}

/* auxiliary: strips last newline from buffer */
void stripn(char *msg)
{
	char *msg_end = msg + strlen(msg);
	while (msg_end >= msg)
	{
		if (*msg_end == '\n' || *msg_end == '\r')
		{
			*msg_end = '\0';
			break;
		}
		msg_end--;
	}
}

/* attempt to connect a socket to its endpoint */
int doconnect(struct connection *new_connection, struct agent *agent, int index)
{

	struct sockaddr_in serv_addr;

	/* init */
	bzero((char *) &serv_addr, sizeof(serv_addr));

	/* set address family to IPv4 */
	serv_addr.sin_family = AF_INET;

	/* set destination address */
	bcopy((char *) agent->host->h_addr, (char *) &serv_addr.sin_addr.s_addr,
			agent->host->h_length);

	/* set destination port */
	serv_addr.sin_port = htons(agent->port);

	/* create an IPv4 TCP socket */
	new_connection->sockfd = socket(AF_INET, SOCK_STREAM, 0);
	if (new_connection->sockfd < 0)
		error("ERROR opening socket");

	/* try to connect */
	if (connect(new_connection->sockfd, (struct sockaddr *) &serv_addr,
			sizeof(serv_addr)) < 0)
	{
		warning("ERROR connecting");
		return 0;
	}

	char buf[16];
	sprintf(buf, "Agent%d", index);

	if (auth_active(new_connection, "forge.pem", "splitter_key.pem", buf) < 0)
	{
		fprintf(stderr, "Auth with Agent failed!\n");
		close(new_connection->sockfd);
		return 0;
	}

	return 1;
}

/* cleanup function, called upon splitter termination */
void terminate(void)
{

	int i;

	for (i = 0; i < config_count; i++)
	{

		close(client_listeners[i]);

		/* close connections with clients */
		close(client_connections[i]->sockfd);

		/* destroy message buffer */
		buf_destroy(client_connections[i]->buffer);

		/* free memory for client connection struct */
		free(client_connections[i]);

		struct connection *connection = agent_connections[i];
		struct connection *tmp;
		while (connection)
		{
			close(connection->sockfd);
			buf_destroy(connection->buffer);
			tmp = connection;
			connection = connection->next;
			free(tmp);
		}
	}

	free(client_listeners);
	free(client_connections);
	free(agent_connections);

	free_config();
}

/* signal handler */
void sighandler(int i)
{
	switch (i)
	{
	case SIGPIPE:
		debug_print("Caught: Broken Pipe\n");
		break;
	case SIGINT:
		fprintf(stderr, "\n%s:%s Caught: SIGINT\n", __FILE__, __func__);
		terminate();
		exit(0);
		break;
	case SIGQUIT:
		fprintf(stderr, "\n%s:%s Caught: SIGQUIT\n", __FILE__, __func__);
		terminate();
		exit(0);
		break;
	case SIGTERM:
		fprintf(stderr, "\n%s:%s Caught: SIGTERM\n", __FILE__, __func__);
		terminate();
		exit(0);
		break;
	default:
		fprintf(stderr, "%s:%s Caught: signal %d\n", __FILE__, __func__, i);
		terminate();
		exit(0);
		break;
	}
}

int main(int argc, char *argv[])
{

	/* auxiliary */
	int i, j, r, n;

	/* holds name of configuration file
	 (read as command-line argument) */
	char *config_filename = NULL;

	/* max file descriptor number
	 used by select(2) */
	unsigned int max_fd = 0;

	/* file descriptor sets.
	 since select(2) modifies the set passed
	 to it as an argument, we need to reset
	 it before each call using a master set. */
	fd_set master_readfd;
	fd_set working_readfd;

	FD_ZERO(&master_readfd);

	/* message buffer. used for reading incoming messages. */
	char buffer[MAX_READ + 1];

	if (argc < 3)
	{
		fprintf(stdout, "Usage: %s -c CONFIG_FILE\n", argv[0]);
		return -1;
	}

	/* parse command-line arguments */
	char c;
	while ((c = getopt(argc, argv, "c:")) != -1)
	{
		switch (c)
		{
		case 'c':
			config_filename = optarg;
			break;
		default:
			info_print("ignoring param %c\n", c);
			break;
		}
	}

	if (!config_filename)
	{
		fprintf(stderr, "error: missing config file\n");
		return -1;
	}

	/* register signal handler for signals */
	signal(SIGPIPE, &sighandler);
	signal(SIGINT, &sighandler);
	signal(SIGQUIT, &sighandler);
	signal(SIGTERM, &sighandler);

	/* read configuration file */
	parse(config_filename);

	/* count active configurations.
	 used to allocate memory for
	 connection structures */
	for (i = 0; i <= MAX_GROUP_N; i++)
	{
		if (agents[i])
			config_count++;
	}

	/* allocate memory for client listeners */
	client_listeners = calloc(config_count, sizeof(int));
	if (!client_listeners)
		error("calloc");

	/* allocate memory for array of pointers to client connections
	 * note: need to explicitly allocate memory for each client
	 * connection structure */
	client_connections = calloc(config_count, sizeof(struct connection *));
	if (!client_connections)
		error("calloc");

	/* allocate memory for array of pointers to agent connections
	 * note: need to explicitly allocate memory for each agent
	 * connection structure */
	agent_connections = calloc(config_count, sizeof(struct connection *));
	if (!agent_connections)
		error("calloc");

	/* prepare client and agent connection structures */
	for (i = 0, j = config_count; i <= MAX_GROUP_N && j >= 0; i++)
	{

		/* prepare agent connection structures.
		 structures for agents in the same configuration
		 are organized in a linked list with the primary
		 agent being at the head of the list. */
		struct agent *agent = agents[i];
		while (agent)
		{

			debug_print("* Processing group %u, preping for agent %s:%u...\n",
					agent->group, agent->host->h_name, agent->port);

			struct connection *new_connection = calloc(1,
					sizeof(struct connection));
			if (!new_connection)
				error("malloc");

			/* pointer to agent (configuration) structure */
			new_connection->agent = agent;

			new_connection->next = NULL;

			new_connection->seq = 1;

			/* get primary connection for group (config_count - j) */
			struct connection *tmp = agent_connections[config_count - j];
			/* reach last node of group */
			while (tmp && tmp->next)
			{
				tmp = tmp->next;
			}

			/* this new connection structure is the primary for this group */
			if (!agent_connections[config_count - j])
				agent_connections[config_count - j] = new_connection;
			/* this new conenction structure is a secondary for this group */
			else
				tmp->next = new_connection;

			/* set up a socket descriptor for this connection
			 * note: we do not connect yet */
			new_connection->sockfd = socket(AF_INET, SOCK_STREAM, 0);
			if (new_connection->sockfd < 0)
				error("ERROR opening socket");

			new_connection->connected = FALSE;

			/* initialize message buffer for this connection
			 note: used for buffering of outgoing messages
			 in case the connection is (temporarily) down */
			new_connection->buffer = buf_init(MSG_BUFFER_SIZE, MAX_READ + 1);
			if (!new_connection->buffer)
				error("buf_init");

			agent = agent->next;
		}

		/* prepare socket client listeners and connections */
		agent = agents[i];
		if (agent)
		{

			debug_print("* Processing group %u, preping for client %u...\n",
					agent->group, config_count - j);

			/* set up a socket descriptor for listening to for this client */
			client_listeners[config_count - j] = socket(AF_INET, SOCK_STREAM,
					0);
			if (client_listeners[config_count - j] < 0)
				error("ERROR opening socket");

			int opt = 1;
			if (setsockopt(client_listeners[config_count - j], SOL_SOCKET,
					SO_REUSEADDR, (char *) &opt, sizeof(opt)) < 0)
				perror("setsockopt");

			struct sockaddr_in serv_addr;

			bzero((char *) &serv_addr, sizeof(serv_addr));

			/* set address family to IPv4 */
			serv_addr.sin_family = AF_INET;

			/* set socket address to ANY
			 (i.e., listen to ANY network interface) */
			serv_addr.sin_addr.s_addr = INADDR_ANY;

			/* set socket port (to listen to) */
			serv_addr.sin_port = htons(agent->group);

			/* bind socket to port */
			if (bind(client_listeners[config_count - j],
					(struct sockaddr *) &serv_addr, sizeof(serv_addr)) < 0)
				error("ERROR on binding");

			/* enter listening mode */
			listen(client_listeners[config_count - j], 1);

			/* add socket descriptor to descriptor set (used by select(2)) */
			FD_SET(client_listeners[config_count-j], &master_readfd);
			if (client_listeners[config_count - j] > max_fd)
				max_fd = client_listeners[config_count - j];

			/* allocate memory for a client connection structure */
			client_connections[config_count - j] = calloc(1,
					sizeof(struct connection));
			if (!client_connections[config_count - j])
				error("calloc");

			/* set connected state to FALSE */
			client_connections[config_count - j]->connected = FALSE;

			client_connections[config_count - j]->seq = 1;

			/* initialize message buffer.
			 messages that fail delivery to the endpoint
			 of this connection, will get buffered for
			 future delivery */
			client_connections[config_count - j]->buffer = buf_init(
					MSG_BUFFER_SIZE, MAX_READ + 1);
			if (!client_connections[config_count - j]->buffer)
				error("buf_init");

			j--;
		}
	}

	/* go go go */
	while (1)
	{

		/* reset working set of file descriptors */
		working_readfd = master_readfd;

		debug_print("- Idle...\n");

		r = select(max_fd + 1, &working_readfd, NULL, NULL, NULL);

		if (r == 0)
			continue; /* select timeout */
		if (r == -1)
		{
			perror("select");
			continue;
		}

		/* check active descriptors (set by select) in working set */
		for (i = 0; i < config_count; i++)
		{

			/* ---------------------------------------------------------- */
			/* check CLIENT listeners (new clients attempting to CONNECT) */
			/* ---------------------------------------------------------- */
			if (FD_ISSET(client_listeners[i], &working_readfd))
			{

				info_print("[C] #%u of client listeners is ready "
				"to accept a new client. Accepting...\n", i);

				struct sockaddr_in cli_addr;
				socklen_t clilen = sizeof(cli_addr);

				/* accept connection, receive socket descriptor for it */
				client_connections[i]->sockfd = accept(client_listeners[i],
						(struct sockaddr *) &cli_addr, &clilen);

				if (client_connections[i]->sockfd < 0)
				{
					perror("accept");
					continue;
				}

				if (auth_passive(client_connections[i], "splitter_cert.pem",
						"splitter_key.pem", "Client", TRUE) < 0)
				{
					fprintf(stderr, "Auth with Client failed\n");
					close(client_connections[i]->sockfd);
					continue;
				}

				/* set connected state to TRUE */
				client_connections[i]->connected = TRUE;

				/* add socket descriptor of this new connection
				 to set monitored by select(2) */
				FD_SET(client_connections[i]->sockfd, &master_readfd);
				if (client_connections[i]->sockfd > max_fd)
					max_fd = client_connections[i]->sockfd;

				/* remove this client's listener descriptor from set
				 so that one instance of each client is active. */
				FD_CLR(client_listeners[i], &master_readfd);
			}

			/* ------------------------------------------------------ */
			/* check CLIENT connections for input (something to READ) */
			/* ------------------------------------------------------ */
			if (client_connections[i]->connected
					&& FD_ISSET(client_connections[i]->sockfd, &working_readfd))
			{

				info_print(
						"[C] #%u of client connections is ready to be read...\n",
						i);

				bzero(buffer, MAX_READ + 1);

				n = sc_read(client_connections[i], buffer, MAX_READ);

				if (n < 0)
					error("ERROR reading from client socket");
				/* client disconnected */
				else if (n == 0)
				{
					info_print(
							"[C] #%u of client connections has DISCONNECTED\n",
							i);

					/* remove this connection's descriptor from set (cleanup) */
					FD_CLR(client_connections[i]->sockfd, &master_readfd);

					/* close descriptor */
					close(client_connections[i]->sockfd);

					/* re-add this client's listener descriptor to set
					 so that future connection attempts are accepted */
					FD_SET(client_listeners[i], &master_readfd);
				}
				else
				{

					info_print(
							"[C] #%u of client connections has read %u bytes\n",
							i, n);
					debug_print(
							"[C] #%u of client connections has read '%s' (%u)\n",
							i, buffer, n);

					/* write buffer (message) to agents.
					 * try to connect (once) on write fail. */
					struct connection *connection = agent_connections[i];
					int agent_index = 1;
					while (connection)
					{

						/* buffer the message just in case delivery fails */
						buf_push2n(connection->buffer, buffer, n);

						/* while buffer has messages, try to transmit them
						 on write fail (e.g. disconnected agent) try once to reconnect */
						char *s = NULL;
						while ((s = buf_peek2n(connection->buffer, &n)))
						{

							if (!connection->connected)
								r = -1;
							else
							{
								r = sc_write(connection, s, n);
							}

							/* write failure */
							if (r <= 0)
							{

								/* remove descriptor of disconnected socket
								 from monitored set */
								FD_CLR(connection->sockfd, &master_readfd);

								info_print(
										"[A] Transmitting to agent %s:%u... %s",
										connection->agent->host->h_name, connection->agent->port, "DISCONNECTED. Reconnecting...\n");

								/* do connect attempt */
								if (doconnect(connection, connection->agent, agent_index))
									connection->connected = TRUE;
								else
									goto secondfail;

								/* try again to write */
								r = sc_write(connection, s, n);

								/* write success */
								if (r > 0)
								{

									info_print(
											"[A] Transmitting to agent %s:%u... "
											"SUCCESS (%d)\n",
											connection->agent->host->h_name, connection->agent->port, r);

									/* pop message from outgoing buffer */
									buf_pop2(connection->buffer);

									/* since write was a success (after a failure),
									 this is a fully functional socket.
									 add its descriptor to
									 set monitored by select(2) */
									FD_SET(connection->sockfd, &master_readfd);
									if (connection->sockfd > max_fd)
										max_fd = connection->sockfd;
								}
								/* write failure (2nd) */
								else
								{
secondfail:							connection->connected = FALSE;

									info_print(
											"[A] Transmitting to agent %s:%u... "
											"FAILURE\n",
											connection->agent->host->h_name, connection->agent->port);

									break;
								}
							}
							/* write success */
							else
							{

								info_print(
										"[A] Transmitting to agent %s:%u... "
										"SUCCESS (%d)\n",
										connection->agent->host->h_name, connection->agent->port, r);

								/* pop message from outgoing buffer */
								buf_pop2(connection->buffer);
							}
						}

						/* move to next connection (agent) for this client */
						connection = connection->next;
						agent_index++;
					}
				}
			}

			/* ----------------------------------------------------- */
			/* check AGENT connections for input (something to READ) */
			/* ignore all input but that coming from primary agent   */
			/* ----------------------------------------------------- */
			struct connection *aconnection = agent_connections[i];
			while (aconnection)
			{

				if (aconnection->connected
						&& FD_ISSET(aconnection->sockfd, &working_readfd))
				{

					info_print(
							"[A] #%u,%u of agent connections is ready to be read..\n",
							i, (unsigned int) ((aconnection - agent_connections[i]) / sizeof(struct connection)));

					bzero(buffer, MAX_READ + 1);

					/* we monitor and try to read from all agent connections
					 to detect disconnected endpoints. we consider only
					 input received from primary agent. */
					n = read(aconnection->sockfd, buffer, MAX_READ);
					if (n < 0)
						error("ERROR reading from client socket");
					/* client disconnected */
					else if (n == 0)
					{

						info_print(
								"[A] #%u,%u of agent connections has DISCONNECTED\n",
								i, (unsigned int) ((aconnection - agent_connections[i]) / sizeof(struct connection)));

						aconnection->connected = FALSE;

						/* remove this connection's descriptor from set (cleanup) */
						FD_CLR(aconnection->sockfd, &master_readfd);

						/* close descriptor */
						close(aconnection->sockfd);
					}
					/* only for primary agent */
					else if (aconnection == agent_connections[i])
					{

						info_print(
								"[A] #%u,%u of agent connections has read %u bytes\n",
								i, (unsigned int) ((aconnection - agent_connections[i]) / sizeof(struct connection)), n);
						debug_print(
								"[A] #%u,%u of agent connections has read '%s' (%u)\n",
								i, (unsigned int) ((aconnection - agent_connections[i]) / sizeof(struct connection)), buffer, n);

						/* write buffer (message) to client.
						 buffer on write fail. */
						struct connection *connection = client_connections[i];

						/* buffer the message just in case delivery fails */
						buf_push2n(connection->buffer, buffer, n);

						/* while buffer has messages, try to transmit them */
						char *s = NULL;
						while ((s = buf_peek2n(connection->buffer, &n)))
						{

							if (!connection->connected)
								r = -1;
							else
								r = send(connection->sockfd, s, n,
										MSG_NOSIGNAL);
							/* write failure */
							if (r <= 0)
							{
								info_print("[C] Transmitting to client #%u... "
								"DISCONNECTED. FAILURE.\n", i);
								break;
							}
							/* write success */
							else
							{

								info_print("[C] Transmitting to client #%u... "
								"SUCCESS (%d)\n", i, r);

								/* pop message from outgoing buffer */
								buf_pop2(connection->buffer);
							}
						}
					}
				}

				/* move on to next configuration
				 (group of agent connections) */
				aconnection = aconnection->next;
			}
		}
	} /* end of go go go */

	return 0;
}
